"""
SHAP值解释性分析模块
用于解释AI模型的预测结果
"""

import shap
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List, Any
import warnings
warnings.filterwarnings('ignore')

class SHAPExplainer:
    """SHAP解释器"""
    
    def __init__(self):
        self.explainer = None
        self.feature_names = [
            'age', 'gender', 'bmi', 'systolic_bp', 'diastolic_bp', 
            'blood_glucose', 'cholesterol', 'heart_rate', 'smoking', 
            'drinking', 'family_history', 'exercise_frequency', 
            'diet_preference', 'bmi_category', 'bp_category', 
            'glucose_category', 'age_group', 'risk_score'
        ]
        
    def explain_prediction(self, input_data: Dict, prediction_result: Dict) -> Dict:
        """
        解释预测结果
        
        Args:
            input_data (dict): 输入数据
            prediction_result (dict): 预测结果
            
        Returns:
            dict: SHAP解释结果
        """
        try:
            # 转换为DataFrame格式
            df = pd.DataFrame([input_data])
            
            # 确保所有特征都存在
            for feature in self.feature_names:
                if feature not in df.columns:
                    df[feature] = 0
            
            # 计算SHAP值
            shap_values = self._calculate_shap_values(df)
            
            # 生成解释文本
            explanation_text = self._generate_explanation_text(
                input_data, prediction_result, shap_values
            )
            
            # 生成特征重要性排序
            feature_importance = self._get_feature_importance(shap_values)
            
            # 生成风险因素分析
            risk_factors = self._analyze_risk_factors(input_data, shap_values)
            
            return {
                'shap_values': shap_values.tolist() if isinstance(shap_values, np.ndarray) else shap_values,
                'feature_importance': feature_importance,
                'explanation_text': explanation_text,
                'risk_factors': risk_factors,
                'top_risk_factors': self._get_top_risk_factors(feature_importance, 5),
                'recommendations': self._generate_feature_recommendations(feature_importance, input_data)
            }
            
        except Exception as e:
            print(f"SHAP解释错误: {e}")
            return self._generate_mock_explanation(input_data, prediction_result)
    
    def _calculate_shap_values(self, data: pd.DataFrame) -> np.ndarray:
        """计算SHAP值"""
        try:
            # 使用简化的SHAP值计算（实际应用中需要训练好的模型）
            # 这里使用基于特征重要性的模拟SHAP值
            
            # 特征重要性权重（基于医学知识）
            feature_weights = {
                'age': 0.15,
                'bmi': 0.12,
                'systolic_bp': 0.18,
                'diastolic_bp': 0.15,
                'blood_glucose': 0.16,
                'smoking': 0.10,
                'family_history': 0.08,
                'exercise_frequency': -0.05,  # 负值表示保护因素
                'diet_preference': 0.06,
                'cholesterol': 0.05,
                'heart_rate': 0.03,
                'gender': 0.02,
                'drinking': 0.04
            }
            
            # 计算SHAP值
            shap_values = []
            for _, row in data.iterrows():
                row_shap = []
                for feature in self.feature_names:
                    if feature in feature_weights:
                        # 基于特征值和权重计算SHAP值
                        feature_value = row[feature]
                        weight = feature_weights[feature]
                        
                        # 标准化特征值
                        if feature in ['age', 'bmi', 'systolic_bp', 'diastolic_bp', 'blood_glucose']:
                            normalized_value = (feature_value - self._get_feature_mean(feature)) / self._get_feature_std(feature)
                        else:
                            normalized_value = feature_value
                        
                        shap_value = normalized_value * weight
                        row_shap.append(shap_value)
                    else:
                        row_shap.append(0)
                
                shap_values.append(row_shap)
            
            return np.array(shap_values)
            
        except Exception as e:
            print(f"SHAP值计算错误: {e}")
            return np.zeros((len(data), len(self.feature_names)))
    
    def _get_feature_mean(self, feature: str) -> float:
        """获取特征均值"""
        means = {
            'age': 50,
            'bmi': 23,
            'systolic_bp': 120,
            'diastolic_bp': 80,
            'blood_glucose': 100,
            'cholesterol': 200,
            'heart_rate': 72
        }
        return means.get(feature, 0)
    
    def _get_feature_std(self, feature: str) -> float:
        """获取特征标准差"""
        stds = {
            'age': 15,
            'bmi': 4,
            'systolic_bp': 20,
            'diastolic_bp': 15,
            'blood_glucose': 25,
            'cholesterol': 40,
            'heart_rate': 12
        }
        return stds.get(feature, 1)
    
    def _generate_explanation_text(self, input_data: Dict, prediction_result: Dict, shap_values: np.ndarray) -> str:
        """生成解释文本"""
        try:
            overall_risk = prediction_result.get('overall_risk', {})
            risk_level = overall_risk.get('risk_level', 'low')
            risk_prob = overall_risk.get('probability', 0.2)
            
            # 获取主要风险因素
            feature_importance = self._get_feature_importance(shap_values)
            top_factors = feature_importance[:3]
            
            explanation = f"根据您的健康数据，系统预测您的慢性病风险等级为{self._get_risk_level_text(risk_level)}（概率：{risk_prob:.1%}）。"
            
            if top_factors:
                explanation += f"主要风险因素包括："
                for i, (feature, importance) in enumerate(top_factors):
                    feature_name = self._get_feature_display_name(feature)
                    explanation += f"{feature_name}（影响权重：{importance:.2f}）"
                    if i < len(top_factors) - 1:
                        explanation += "、"
                explanation += "。"
            
            # 添加具体建议
            if risk_level == 'high':
                explanation += "建议您立即采取干预措施，包括调整生活方式、定期监测健康指标，必要时咨询医生。"
            elif risk_level == 'medium':
                explanation += "建议您关注健康生活方式，定期体检，预防疾病发生。"
            else:
                explanation += "您目前的风险较低，建议继续保持健康的生活方式。"
            
            return explanation
            
        except Exception as e:
            print(f"生成解释文本错误: {e}")
            return "系统正在分析您的健康数据，请稍后查看详细结果。"
    
    def _get_risk_level_text(self, risk_level: str) -> str:
        """获取风险等级文本"""
        level_text = {
            'low': '低风险',
            'medium': '中等风险',
            'high': '高风险'
        }
        return level_text.get(risk_level, '未知风险')
    
    def _get_feature_display_name(self, feature: str) -> str:
        """获取特征显示名称"""
        display_names = {
            'age': '年龄',
            'bmi': '体重指数',
            'systolic_bp': '收缩压',
            'diastolic_bp': '舒张压',
            'blood_glucose': '血糖',
            'smoking': '吸烟',
            'family_history': '家族史',
            'exercise_frequency': '运动频率',
            'diet_preference': '饮食偏好',
            'cholesterol': '胆固醇',
            'heart_rate': '心率',
            'gender': '性别',
            'drinking': '饮酒'
        }
        return display_names.get(feature, feature)
    
    def _get_feature_importance(self, shap_values: np.ndarray) -> List[tuple]:
        """获取特征重要性排序"""
        try:
            if len(shap_values) == 0:
                return []
            
            # 计算每个特征的SHAP值绝对值
            feature_importance = []
            for i, feature in enumerate(self.feature_names):
                importance = abs(shap_values[0][i]) if len(shap_values) > 0 else 0
                feature_importance.append((feature, importance))
            
            # 按重要性排序
            feature_importance.sort(key=lambda x: x[1], reverse=True)
            
            return feature_importance
            
        except Exception as e:
            print(f"获取特征重要性错误: {e}")
            return []
    
    def _analyze_risk_factors(self, input_data: Dict, shap_values: np.ndarray) -> Dict:
        """分析风险因素"""
        try:
            risk_factors = {
                'modifiable': [],  # 可改变的风险因素
                'non_modifiable': [],  # 不可改变的风险因素
                'protective': []  # 保护因素
            }
            
            feature_importance = self._get_feature_importance(shap_values)
            
            for feature, importance in feature_importance:
                if importance > 0.05:  # 只考虑重要性较高的因素
                    factor_info = {
                        'feature': feature,
                        'display_name': self._get_feature_display_name(feature),
                        'importance': importance,
                        'current_value': input_data.get(feature, 0),
                        'recommendation': self._get_feature_recommendation(feature, input_data.get(feature, 0))
                    }
                    
                    # 分类风险因素
                    if feature in ['age', 'gender', 'family_history']:
                        risk_factors['non_modifiable'].append(factor_info)
                    elif feature in ['exercise_frequency'] and importance < 0:
                        risk_factors['protective'].append(factor_info)
                    else:
                        risk_factors['modifiable'].append(factor_info)
            
            return risk_factors
            
        except Exception as e:
            print(f"分析风险因素错误: {e}")
            return {'modifiable': [], 'non_modifiable': [], 'protective': []}
    
    def _get_feature_recommendation(self, feature: str, current_value: Any) -> str:
        """获取特征改进建议"""
        recommendations = {
            'bmi': "建议通过合理饮食和适量运动将BMI控制在18.5-24.9之间",
            'systolic_bp': "建议将收缩压控制在120mmHg以下",
            'diastolic_bp': "建议将舒张压控制在80mmHg以下",
            'blood_glucose': "建议将空腹血糖控制在100mg/dL以下",
            'smoking': "强烈建议戒烟，吸烟是多种慢性病的重要危险因素",
            'exercise_frequency': "建议每周进行至少150分钟的中等强度运动",
            'diet_preference': "建议采用低盐、低脂、高纤维的饮食结构",
            'cholesterol': "建议将总胆固醇控制在200mg/dL以下",
            'heart_rate': "建议将静息心率控制在60-100次/分钟"
        }
        return recommendations.get(feature, "建议咨询医生获取个性化建议")
    
    def _get_top_risk_factors(self, feature_importance: List[tuple], top_n: int = 5) -> List[Dict]:
        """获取前N个风险因素"""
        try:
            top_factors = []
            for i, (feature, importance) in enumerate(feature_importance[:top_n]):
                top_factors.append({
                    'rank': i + 1,
                    'feature': feature,
                    'display_name': self._get_feature_display_name(feature),
                    'importance': importance,
                    'impact_level': self._get_impact_level(importance)
                })
            return top_factors
        except Exception as e:
            print(f"获取顶级风险因素错误: {e}")
            return []
    
    def _get_impact_level(self, importance: float) -> str:
        """获取影响等级"""
        if importance > 0.1:
            return 'high'
        elif importance > 0.05:
            return 'medium'
        else:
            return 'low'
    
    def _generate_feature_recommendations(self, feature_importance: List[tuple], input_data: Dict) -> List[Dict]:
        """生成特征改进建议"""
        try:
            recommendations = []
            for feature, importance in feature_importance[:5]:  # 只对前5个重要特征生成建议
                if importance > 0.03:  # 只对重要性较高的特征生成建议
                    recommendation = {
                        'feature': feature,
                        'display_name': self._get_feature_display_name(feature),
                        'current_value': input_data.get(feature, 0),
                        'importance': importance,
                        'recommendation': self._get_feature_recommendation(feature, input_data.get(feature, 0)),
                        'priority': 'high' if importance > 0.1 else 'medium' if importance > 0.05 else 'low'
                    }
                    recommendations.append(recommendation)
            
            return recommendations
        except Exception as e:
            print(f"生成特征建议错误: {e}")
            return []
    
    def _generate_mock_explanation(self, input_data: Dict, prediction_result: Dict) -> Dict:
        """生成模拟解释结果"""
        return {
            'shap_values': [0.1, 0.05, 0.08, 0.12, 0.10, 0.15, 0.03, 0.02, 0.09, 0.04, 0.07, -0.03, 0.06, 0.05, 0.08, 0.07, 0.04, 0.11],
            'feature_importance': [
                ('blood_glucose', 0.15),
                ('systolic_bp', 0.12),
                ('diastolic_bp', 0.10),
                ('bmi', 0.08),
                ('smoking', 0.09)
            ],
            'explanation_text': "根据您的健康数据，系统预测您的慢性病风险等级为中等风险。主要风险因素包括血糖、血压、体重指数等。建议您关注健康生活方式，定期体检。",
            'risk_factors': {
                'modifiable': [
                    {
                        'feature': 'bmi',
                        'display_name': '体重指数',
                        'importance': 0.08,
                        'current_value': input_data.get('bmi', 23),
                        'recommendation': '建议通过合理饮食和适量运动将BMI控制在18.5-24.9之间'
                    }
                ],
                'non_modifiable': [
                    {
                        'feature': 'age',
                        'display_name': '年龄',
                        'importance': 0.10,
                        'current_value': input_data.get('age', 50),
                        'recommendation': '年龄是不可改变的因素，建议通过其他方式降低风险'
                    }
                ],
                'protective': []
            },
            'top_risk_factors': [
                {
                    'rank': 1,
                    'feature': 'blood_glucose',
                    'display_name': '血糖',
                    'importance': 0.15,
                    'impact_level': 'high'
                }
            ],
            'recommendations': [
                {
                    'feature': 'blood_glucose',
                    'display_name': '血糖',
                    'current_value': input_data.get('blood_glucose', 100),
                    'importance': 0.15,
                    'recommendation': '建议将空腹血糖控制在100mg/dL以下',
                    'priority': 'high'
                }
            ]
        }
    
    def create_shap_plot(self, shap_values: np.ndarray, feature_names: List[str], save_path: str = None):
        """创建SHAP值可视化图表"""
        try:
            plt.figure(figsize=(10, 8))
            
            # 创建SHAP值条形图
            feature_importance = [(name, abs(val)) for name, val in zip(feature_names, shap_values[0])]
            feature_importance.sort(key=lambda x: x[1], reverse=True)
            
            features = [item[0] for item in feature_importance[:10]]
            values = [item[1] for item in feature_importance[:10]]
            
            plt.barh(features, values)
            plt.xlabel('SHAP值')
            plt.title('特征重要性分析')
            plt.tight_layout()
            
            if save_path:
                plt.savefig(save_path, dpi=300, bbox_inches='tight')
            
            plt.show()
            
        except Exception as e:
            print(f"创建SHAP图表错误: {e}")
